#!/usr/bin/env python
"""
classify.py is an out-of-the-box image classifer callable from the command line.

By default it configures and runs the Caffe reference ImageNet model.
"""
import numpy as np
import os
import sys
import argparse
import glob
import time
import cv2

import pandas as pd
from skimage.color import rgb2gray

import caffe


global x
global y
x=[]
y=[]


def get_max(prediction):
	max=0
	for i in range(len(prediction)):
		if(prediction[i]>prediction[max]):
			max=i
	return max

def get_max_class(images,predictions,height,width):
	plane=0
	for i in range(len(images)):
		for j in range(len(predictions[i])):
			max=get_max(predictions[i][j])
			if max==0:
				plane+=1
				tmp=images[i][x[j]:x[j]+height,y[j]:y[j]+height]
				cv2.rectangle(images[i],(x[j],y[j]),(x[j]+width,y[j]+height),(0,0,255),2)
		cv2.imwrite('/home/wyj/caffe/examples/eagLeNet/a.jpg',images[i])
		cv2.waitKey(0)
	return plane

def cut_image(image,height,width):
	del x[:]
	del y[:]
	rows,cols,channels=image.shape
	i=0
	imgs=[]
	while i+height<rows:
		j=0
		while j+width<cols:
			img=image[i:i+height,j:j+width]
			imgs.append(img)
			x.append(j)
			y.append(i)
			j+=width/2
		i+=height/2
	return imgs

def main(argv):
    pycaffe_dir = os.path.dirname(__file__)

    parser = argparse.ArgumentParser()
    # Required arguments: input and output files.
    parser.add_argument(
        "input_file",
        help="Input image, directory, or npy."
    )
    parser.add_argument(
        "output_file",
        help="Output npy filename."
    )
    # Optional arguments.
    parser.add_argument(
        "--model_def",
        default=os.path.join(pycaffe_dir,
                "../models/bvlc_reference_caffenet/deploy.prototxt"),
        help="Model definition file."
    )
    parser.add_argument(
        "--pretrained_model",
        default=os.path.join(pycaffe_dir,
                "../models/bvlc_reference_caffenet/bvlc_reference_caffenet.caffemodel"),
        help="Trained model weights file."
    )
    parser.add_argument(
        "--gpu",
        action='store_true',
        help="Switch for gpu computation."
    )
    parser.add_argument(
        "--center_only",
        action='store_true',
        help="Switch for prediction from center crop alone instead of " +
             "averaging predictions across crops (default)."
    )
    parser.add_argument(
        "--images_dim",
        default='256,256',
        help="Canonical 'height,width' dimensions of input images."
    )
    parser.add_argument(
        "--mean_file",
        default=os.path.join(pycaffe_dir,
                             'caffe/imagenet/ilsvrc_2012_mean.npy'),
        help="Data set image mean of [Channels x Height x Width] dimensions " +
             "(numpy array). Set to '' for no mean subtraction."
    )
    parser.add_argument(
        "--input_scale",
        type=float,
        help="Multiply input features by this scale to finish preprocessing."
    )
    parser.add_argument(
        "--raw_scale",
        type=float,
        default=255.0,
        help="Multiply raw input by this scale before preprocessing."
    )
    parser.add_argument(
        "--channel_swap",
        default='2,1,0',
        help="Order to permute input channels. The default converts " +
             "RGB -> BGR since BGR is the Caffe default by way of OpenCV."
    )
    parser.add_argument(
        "--ext",
        default='jpg',
        help="Image file extension to take as input when a directory " +
             "is given as the input file."
    )
    
    # add by caisenchuan
    parser.add_argument(
        "--labels_file",
        default=os.path.join(pycaffe_dir,
                "../data/ilsvrc12/synset_words.txt"),
        help="Readable label definition file."
    )
    parser.add_argument(
        "--print_results",
        action='store_true',
        help="Write output text to stdout rather than serializing to a file."
    )
    parser.add_argument(
        "--force_grayscale",
        action='store_true',
        help="Converts RGB images down to single-channel grayscale versions," +
             "useful for single-channel networks like MNIST."
    )
    
    args = parser.parse_args()

    image_dims = [int(s) for s in args.images_dim.split(',')]

    mean, channel_swap = None, None
    
    # add by caisenchuan
    if args.force_grayscale:
      channel_swap = None
      mean = None
    else:
        if args.mean_file:
            mean = np.load(args.mean_file)
        if args.channel_swap:
            channel_swap = [int(s) for s in args.channel_swap.split(',')]
    
    if args.gpu:
        caffe.set_mode_gpu()
        print("GPU mode")
    else:
        caffe.set_mode_cpu()
        print("CPU mode")

    # Make classifier.
    classifier = caffe.Classifier(args.model_def, args.pretrained_model,
            image_dims=image_dims, mean=mean,
            input_scale=args.input_scale, raw_scale=args.raw_scale,
            channel_swap=channel_swap)

    # Load numpy array (.npy), directory glob (*.jpg), or image file.
    args.input_file = os.path.expanduser(args.input_file)
    if args.input_file.endswith('npy'):
        print("Loading file: %s" % args.input_file)
        inputs = np.load(args.input_file)
    elif os.path.isdir(args.input_file):
        print("Loading folder: %s" % args.input_file)
        images=[cv2.imread(im_f)
        for im_f in glob.glob(args.input_file+'/*.'+args.ext)]
        inputs =[caffe.io.load_image(im_f)
                 for im_f in glob.glob(args.input_file + '/*.' + args.ext)]
    else:
        print("Loading file: %s" % args.input_file)
        inputs = [caffe.io.load_image(args.input_file)]
        images=[cv2.imread(args.input_file)]
    
    if args.force_grayscale:
      inputs = [rgb2gray(input) for input in inputs];
    print '=============='
    print("Classifying %d inputs." % len(inputs))

    # Classify.
    start = time.time()
    predictions = []
    for input in inputs:
    	imgs=cut_image(input,128,128)
    	print(len(imgs))
    	prediction=classifier.predict(imgs,not args.center_only)
    	predictions.append(prediction)
	print '=============='
    	print("plane:%d."%get_max_class(images,predictions,128,128))
    print("Done in %.2f s." % (time.time() - start))
    #print("Predictions : %s" % predictions)

    eag_items=['airplane','space','else']

    # for i in range(10):
    #     print '#',cifar10_items[i],':',predictions[0][i]
    #for i in range(len(eag_items)):
      #  print '#',eag_items[i],':',predictions[0][i]

    # print result, add by caisenchuan
    #if args.print_results:
      #  scores = predictions.flatten()
        # with open(args.labels_file) as f:
        #     labels_df = pd.DataFrame([
        #             {
        #                 'synset_id': l.strip().split(' ')[0],
        #                 'name': ' '.join(l.strip().split(' ')[1:]).split(',')[0]
        #             }
        #             for l in f.readlines()
        #         ])
        #     labels = labels_df.sort('synset_id')['name'].values

        #     indices = (-scores).argsort()[:5]
        #     ps = labels[indices]

        #     meta = [
        #         (p, '%.5f' % scores[i])
        #         for i, p in zip(indices, ps)
        #     ]

        #     print meta

    print '=============='
    # Save
    print("Saving results into %s" % args.output_file)
    np.save(args.output_file, predictions)


if __name__ == '__main__':
    main(sys.argv)
